// ProvHelperNative.cpp
// This file demos the way you encrypt a prov.xml that can be accepted by WEH8.1 Oobe
//

#include "EncryptProv.h"

void Usage(int ret)
{
    wprintf(L"Usage: EncryptProv [source file] [target file] [password]\r\n");
    wprintf(L"Usage: Accepted password length is from 1 to 260");
    wprintf(L"Usage: Accepted password chars are normal ASCII chars, such as: A-Za-z0-9,./;'[]\-=<>?, etc..");
    exit(ret);
}

// acquire a handle to a particular key container within a particular cryptographic service provider (CSP). 
// This returned handle is used in calls to CryptoAPI functions that use the selected CSP.
// http://msdn.microsoft.com/en-us/library/windows/desktop/aa379886(v=vs.85).aspx
HRESULT AcquireContext(
    DWORD nAlgorithm,
    HCRYPTPROV& hProv)
{
    HRESULT hr = S_OK;

    if (nAlgorithm == CALG_AES_128 || nAlgorithm == CALG_AES_192 || nAlgorithm == CALG_AES_256)
    {
        CHECKWIN32 (CryptAcquireContext (&hProv, NULL, MS_ENH_RSA_AES_PROV, PROV_RSA_AES, CRYPT_VERIFYCONTEXT));
    }
    else
    {
        CHECKWIN32 (CryptAcquireContext (&hProv, NULL, NULL, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT));
    }

Exit:

    return hr;
}

// internal function to create hash key, returns HCRYPTKEY for later actual encrytion
HRESULT CreateHashKey(
    HCRYPTPROV hProv,
    PBYTE pbKey,
    DWORD nKeySize,
    DWORD nHashAlgorithm,
    DWORD nEncryptAlgorithm, 
    HCRYPTKEY& hKey        
    )
{
    HRESULT hr = S_OK;
    HCRYPTHASH hHash = NULL;
    CHECKPARAM(hProv);

    // initiates the hashing of a stream of data
    CHECKWIN32 (CryptCreateHash (hProv, nHashAlgorithm, 0, 0, &hHash));

    // adds data to a specified hash object
    CHECKWIN32 (CryptHashData (hHash, pbKey, nKeySize, 0));

    // generates cryptographic session keys derived from a base data value
    CHECKWIN32 (CryptDeriveKey (hProv, nEncryptAlgorithm, hHash, 0, &hKey));

Exit:
    if (hHash != NULL) { CryptDestroyHash (hHash); }

    return hr;
}

// internal function to encrypt data
HRESULT EncryptBufferNoHeader (
    PBYTE  pbKey,      
    DWORD  nKeySize,   
    PBYTE  pbInput,    
    DWORD  nInputSize, 
    PBYTE &pbEncrypted,
    DWORD &nEncryptedSize,
    DWORD  nAlgorithm
    )
{
    HRESULT hr = S_OK;
    HCRYPTPROV hProv = NULL;
    HCRYPTKEY  hKey  = NULL;

    PBYTE pData = NULL;

    CHECKPARAM (pbKey);
    CHECKPARAM (pbInput);

    pbEncrypted = NULL;
    nEncryptedSize = 0;

    // see comments inside this function
    CHECKHRRET(AcquireContext(nAlgorithm, hProv));

    // see comments inside this function
    CHECKHRRET(CreateHashKey(
        hProv,
        pbKey,
        nKeySize,
        CALG_SHA_256,
        nAlgorithm,
        hKey));

    DWORD dwBlockLen = 0;
    DWORD dwLen = sizeof (DWORD);

    // retrieves data that governs the operations of a key. 
    // If the Microsoft Cryptographic Service Provider is used, the base symmetric keying material is not obtainable by this or any other function.
    CHECKWIN32 (CryptGetKeyParam (hKey, KP_BLOCKLEN, (BYTE*) &dwBlockLen, &dwLen, 0));

    DWORD nBufferSize = nInputSize + ((dwBlockLen + 7) / 8);

    pData = new BYTE [nBufferSize];
    CHECKMEM (pData);

    CopyMemory (pData, pbInput, nInputSize);

    // encrypts data
    CHECKWIN32 (CryptEncrypt (hKey, 0, TRUE, 0, pData, &nInputSize, nBufferSize));

    nEncryptedSize = nInputSize;
    pbEncrypted = pData;
    pData = NULL;

Exit:

    if (pData != NULL) { delete []pData; }

    if (hKey  != NULL) { CryptDestroyKey (hKey); }
    if (hProv != NULL) { CryptReleaseContext (hProv, 0); }

    return hr;
}

// internal function to encrypt and encode data
HRESULT EncryptAndEncodeBuffer (
    PBYTE  pbKey,     
    DWORD  nKeySize,       
    PBYTE  pbInput,        
    DWORD  nInputSize,     
    PBYTE &pbOutput,       
    DWORD &nOutputSize,    
    DWORD  nAlgorithm
    )
{
    HRESULT hr = S_OK;
    PBYTE pbEncrypted = NULL;
    DWORD nEncryptedSize;
    int nEncodedSize;

    // perform encryption
    CHECKHRRET(EncryptBufferNoHeader(
        pbKey,
        nKeySize,
        pbInput,
        nInputSize,
        pbEncrypted,
        nEncryptedSize,
        nAlgorithm));

    // get required length to allocate memory for encode
    nEncodedSize = Base64EncodeGetRequiredLength(nEncryptedSize);

    pbOutput = new BYTE[nEncodedSize];
    CHECKMEM (pbOutput);

    // perform BASE64 encode
    CHECK(Base64Encode(pbEncrypted, nEncryptedSize, (LPSTR)pbOutput, &nEncodedSize), E_FAIL);
    nOutputSize = nEncodedSize;

Exit:
    if (pbEncrypted != NULL) { delete []pbEncrypted; }

    return hr;
}

int _tmain(int argc, _TCHAR* argv[])
{	
    CComPtr<IXMLDOMDocument> m_pProvXmlDocument;
    CComPtr<IXMLDOMParseError> pError;
    CComBSTR sReason, sSource;
    long nLine = 0, nColumn = 0;
    HRESULT     hr = S_OK;
    PBYTE       password = NULL;
    WCHAR       *pbFileOriginal = NULL;
    PBYTE       pbFileFinal = NULL;
    BSTR        xmlString;
    DWORD       fileSizeOriginal;
    DWORD       fileSizeFinal;
    size_t      cchPassword = 0;
    HANDLE      hFileTarget = INVALID_HANDLE_VALUE;
    VARIANT_BOOL isSuccess;
    CComVariant sourceFilePath(argv[1]);

    if (argc != 4)
    {
        wprintf(L"Invalid arguements");
        Usage(1);
    }

    password = (PBYTE) argv[3];

    cchPassword = wcsnlen(argv[3], MAX_PASSWORD);

    CHECKHR(CoInitializeEx(NULL, COINITBASE_MULTITHREADED));

    CHECKHR(m_pProvXmlDocument.CoCreateInstance(__uuidof(DOMDocument60), nullptr, CLSCTX_INPROC_SERVER));

    // load the source file as XML and leverage IXMLDOMDOCUMENT built in methods to validate the source file
    // and print any validation errors
    // also this way naturally handles source file encoding difference, such as ANSI, UTF8, UTF16, etc.., eventually it will be WCHAR string before encryption
    hr = m_pProvXmlDocument->load(sourceFilePath, &isSuccess);

    if(S_OK != hr)
    {
        m_pProvXmlDocument->get_parseError(&pError);

        if(pError)
        {
            pError->get_reason(&sReason);
            pError->get_srcText(&sSource);
            pError->get_line(&nLine);
            pError->get_linepos(&nColumn);

            wprintf(L"XML parse error was detected\r\n");
            wprintf(L"Reason:   %s\r\n", sReason);
            wprintf(L"Source:   %s\r\n", sSource);
            wprintf(L"Line:     %d\r\n", nLine);
            wprintf(L"Column:   %d\r\n", nColumn);
        }

        goto Exit;
    }

    CHECKHR(m_pProvXmlDocument->get_xml(&xmlString));

    fileSizeOriginal = SysStringLen(xmlString) + 1;

    pbFileOriginal = new WCHAR[fileSizeOriginal];

    if (NULL == pbFileOriginal)
    {
        hr = E_OUTOFMEMORY;
        goto Exit;
    }

    // copy the xml content into a WCHAR string, which is the required format for Oobe to decrypt
    if(0 != wcsncpy_s(pbFileOriginal, fileSizeOriginal, xmlString, _TRUNCATE))
    {
        goto Exit;
    }

    // perform the encryption and BASE64 encode
    hr = EncryptAndEncodeBuffer(
        password, 
        cchPassword * sizeof(WCHAR), 
        (PBYTE)pbFileOriginal, 
        fileSizeOriginal * sizeof(WCHAR),
        pbFileFinal, 
        fileSizeFinal, 
        CALG_AES_128);

    if (FAILED(hr))
    {
        goto Exit;
    }

    hFileTarget = CreateFile(
        argv[2],
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_WRITE,
        NULL,
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,
        NULL);

    if (FALSE == WriteFile(hFileTarget, pbFileFinal, fileSizeFinal, NULL, NULL))
    {
        hr = HRESULT_FROM_WIN32(GetLastError());
        goto Exit;
    }

    if (FAILED(SetEndOfFile(hFileTarget)))
    {
        hr = HRESULT_FROM_WIN32(GetLastError());
        goto Exit;
    }

Exit:

    m_pProvXmlDocument.Release();

    pError.Release();

    CoUninitialize();

    if(hFileTarget != INVALID_HANDLE_VALUE)
    {
        CloseHandle(hFileTarget);
    }
    if(pbFileFinal)
    {
        delete []pbFileFinal;
    }
    if(pbFileOriginal)
    {
        delete []pbFileOriginal;
    }

    return (int)hr;
}